import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { Chroma } from "@langchain/community/vectorstores/chroma";
import { PDFLoader } from "langchain/document_loaders/fs/pdf";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { ChromaClient } from 'chromadb'

const splitPDF =async()=>{
    const loader = new PDFLoader("./doc/愚公移山.pdf");
    const docs = await loader.load();
    console.log(docs.length)
    // console.log(docs[0])
    const textSplitter = new RecursiveCharacterTextSplitter({
        chunkSize: 50,
        chunkOverlap: 10,
      });
    const splitDocs = await textSplitter.splitDocuments(docs);
    // console.log({ splitDocs });
    console.log(splitDocs.length)
    const embeddings = new OllamaEmbeddings({
        model: "nomic-embed-text", // default value
        baseUrl: "http://117.72.38.226:11434", // default value
        requestOptions: {
        // useMMap: true,
        numThread: 6,
        // numGpu: 1,
        },
    });
    const vectorStore= await Chroma.fromDocuments(splitDocs, embeddings, {
        collectionName: "test-hcy-3", // id  14486ba6-2d74-4658-a3f0-0a48a7e3455d
        url: "http://117.72.38.226:8000", // Optional, will default to this value
    });
    const response =await vectorStore.similaritySearch("山神做了什么？",3)
    console.log(response)
}

const search =async()=>{
    const embeddings = new OllamaEmbeddings({
        model: "nomic-embed-text", // default value
        baseUrl: "http://117.72.38.226:11434", // default value
        requestOptions: {
        // useMMap: true,
        numThread: 6,
        // numGpu: 1,
        },
    });
    const vectorStore = await Chroma.fromExistingCollection(
        embeddings,
        {
            collectionName: "test-hcy-3",
            url: "http://117.72.38.226:8000",
        }
      );
    const response =await vectorStore.similaritySearch("山神做了什么？",3)
    console.log(response)
}

const clearChroma =async ()=>{
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    await client.deleteCollection({
        name:"test-hcy-3"
    })
    const list =await client.listCollections()
    console.log(list)
    // const coll =await client.getCollection({
    //             name: "test-hcy-3",
    //           });
    // const datas =await coll.get()
    // console.log(datas)
    // for (let i = 0; i < list.length; i++) {
    //     const coll =await client.getCollection({
    //         name: list[i].name,
    //       });
    //     const li=await coll.count()
    //     console.log(list[i].name,li)
    // }
}


splitPDF()
// search()
// clearChroma()